from functools import cached_property
import logging
from pathlib import Path
import polars as pl
import cnvpytor
from pydantic import BaseModel
from .utils import is_active_chrom

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("ClinCNVseq")


class MergeCN(BaseModel):
    pytor_file: Path
    bin_size: int
    gap_file: Path
    mask_cn_file: Path
    raw_cn_file: Path
    multi_cn_file: Path
    control_cn_files: list[Path]
    gender: str

    class Config:
        arbitrary_types_allowed = True

    @cached_property
    def view(self) -> cnvpytor.Viewer:
        return cnvpytor.Viewer([str(self.pytor_file)], params={"bin_size": self.bin_size})

    @cached_property
    def gap_records(self) -> pl.DataFrame:
        return pl.read_csv(self.gap_file, separator="\t", has_header=False, columns=[1, 2, 3], new_columns=["chromosome", "start", "end"])

    @cached_property
    def chromosomes(self) -> list[str]:
        ref_genome = self.view.reference_genome
        if not ref_genome:
            raise ValueError("Error: view.reference_genome is None")
        return list(filter(lambda x: is_active_chrom(x, self.gender), ref_genome["chromosomes"].keys()))

    def calc_cn_records(self, use_mask: bool) -> pl.DataFrame:
        logger.info(f"Calling read depths with bin size: {self.bin_size} and use_mask: {use_mask}")
        mask_flag = cnvpytor.FLAG_GC_CORR | cnvpytor.FLAG_USEMASK
        raw_flag = cnvpytor.FLAG_GC_CORR
        flag = mask_flag if use_mask else raw_flag
        mean = self.view.io[0].rd_normal_level(self.bin_size, flag)[0]
        rows = []
        for chrom in self.chromosomes:
            rd = pl.Series(self.view.io[0].get_signal(chrom, self.bin_size, "RD", flag))
            for i, cn in enumerate(2 * rd / mean):
                if not isinstance(cn, (float, int)):
                    continue
                start = self.bin_size * i + 1
                end = self.bin_size * (i + 1)
                in_gap = not self.gap_records.filter(pl.col("chromosome").eq(chrom) & pl.col("start").le(end) & pl.col("end").ge(start)).is_empty()
                if str(cn) == "nan" or (not use_mask and cn == 0 and in_gap):
                    continue
                rows.append([chrom, start, end, cn])
        return pl.DataFrame(rows, schema={"chromosome": str, "start": int, "end": int, "cn": float})

    def do_cn(self):
        if not self.mask_cn_file.exists():
            self.calc_cn_records(use_mask=True).write_csv(self.mask_cn_file, separator="\t")
        if not self.raw_cn_file.exists():
            self.calc_cn_records(use_mask=False).write_csv(self.raw_cn_file, separator="\t")

    def do_merge_cn(self):
        logger.info("Merge copy number value of multi-samples to %s", self.multi_cn_file)
        records = pl.DataFrame()
        for i, cn_file in enumerate([self.raw_cn_file] + self.control_cn_files):
            sample = cn_file.name.split(".")[0]
            t_records = pl.read_csv(cn_file, separator="\t")
            t_records = t_records.select(["chromosome", "start", "end", "cn"]).rename({"cn": sample})
            records = records.join(t_records, on=["chromosome", "start", "end"], how="left") if i > 0 else t_records
        records.write_csv(self.multi_cn_file, separator="\t")

    def run(self):
        if not self.mask_cn_file.exists() or not self.raw_cn_file.exists():
            self.do_cn()
        if not self.multi_cn_file.exists():
            self.do_merge_cn()
